import time
import sys
from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_
from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_
from unitree_sdk2py.g1.loco.g1_loco_client import LocoClient
import math
from dataclasses import dataclass

@dataclass
class TestOption:
    name: str  # 测试选项的名称
    id: int    # 测试选项的唯一标识符

# 定义测试选项列表，每个选项对应一个机器人动作
option_list = [
    TestOption(name="damp", id=0),         # 阻尼模式
    TestOption(name="Squat2StandUp", id=1),     # 从蹲下到站立
    TestOption(name="StandUp2Squat", id=2),     # 从站立到蹲下
    TestOption(name="move forward", id=3),      # 向前移动
    TestOption(name="move lateral", id=4),      # 横向移动
    TestOption(name="move rotate", id=5),       # 原地旋转
    TestOption(name="low stand", id=6),         # 低姿态站立
    TestOption(name="high stand", id=7),        # 高姿态站立
    TestOption(name="zero torque", id=8),       # 零力矩模式
    TestOption(name="wave hand1", id=9),        # 挥手（不转身）
    TestOption(name="wave hand2", id=10),       # 挥手并转身
    TestOption(name="shake hand", id=11),       # 摇手
    TestOption(name="Lie2StandUp", id=12),      # 从躺下到站立
]

class UserInterface:
    def __init__(self):
        self.test_option_ = None  # 当前选择的测试选项

    def convert_to_int(self, input_str):
        # 尝试将输入字符串转换为整数
        try:
            return int(input_str)
        except ValueError:
            return None

    def terminal_handle(self):
        # 处理用户输入
        input_str = input("Enter id or name: \n")  # 提示用户输入测试选项的名称或ID

        if input_str == "list":
            # 如果输入 "list"，列出所有测试选项
            self.test_option_.name = None
            self.test_option_.id = None
            for option in option_list:
                print(f"{option.name}, id: {option.id}")
            return

        # 匹配用户输入的名称或ID
        for option in option_list:
            if input_str == option.name or self.convert_to_int(input_str) == option.id:
                self.test_option_.name = option.name
                self.test_option_.id = option.id
                print(f"Test: {self.test_option_.name}, test_id: {self.test_option_.id}")
                return

        # 如果没有匹配的选项，提示用户
        print("No matching test option found.")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        # 如果没有提供网络接口参数，提示用户正确使用
        print(f"Usage: python3 {sys.argv[0]} networkInterface")
        sys.exit(-1)

    print("WARNING: Please ensure there are no obstacles around the robot while running this example.")
    input("Press Enter to continue...")  # 提示用户确认继续

    # 初始化通信通道
    ChannelFactoryInitialize(0, sys.argv[1])

    # 初始化测试选项和用户界面
    test_option = TestOption(name=None, id=None) 
    user_interface = UserInterface()
    user_interface.test_option_ = test_option

    # 初始化运动客户端
    sport_client = LocoClient()  
    sport_client.SetTimeout(10.0)  # 设置超时时间
    sport_client.Init()  # 初始化客户端

    print("Input \"list\" to list all test option ...")  # 提示用户输入 "list" 查看所有选项
    while True:
        user_interface.terminal_handle()  # 处理用户输入

        print(f"Updated Test Option: Name = {test_option.name}, ID = {test_option.id}")

        # 根据用户选择的测试选项执行对应的动作
        if test_option.id == 0:
            sport_client.Damp()  # 阻尼模式
        elif test_option.id == 1:
            sport_client.Damp()
            time.sleep(0.5)
            sport_client.Squat2StandUp()  # 从蹲下到站立
        elif test_option.id == 2:
            sport_client.StandUp2Squat()  # 从站立到蹲下
        elif test_option.id == 3:
            sport_client.Move(0.3, 0, 0)  # 向前移动
        elif test_option.id == 4:
            sport_client.Move(0, 0.3, 0)  # 横向移动
        elif test_option.id == 5:
            sport_client.Move(0, 0, 0.3)  # 原地旋转
        elif test_option.id == 6:
            sport_client.LowStand()  # 低姿态站立
        elif test_option.id == 7:
            sport_client.HighStand()  # 高姿态站立
        elif test_option.id == 8:
            sport_client.ZeroTorque()  # 零力矩模式
        elif test_option.id == 9:
            sport_client.WaveHand()  # 挥手（不转身）
        elif test_option.id == 10:
            sport_client.WaveHand(True)  # 挥手并转身
        elif test_option.id == 11:
            sport_client.ShakeHand()  # 摇手
            time.sleep(3)
            sport_client.ShakeHand()
        elif test_option.id == 12:
            sport_client.Damp()
            time.sleep(0.5)
            sport_client.Lie2StandUp()  # 从躺下到站立（确保机器人面朝上，地面平坦且粗糙）

        time.sleep(1)  # 每次动作之间等待1秒